#!/usr/bin/env python3
import rospy
from geometry_msgs.msg import Twist
from std_msgs.msg import Float64MultiArray
from sensor_msgs.msg import JointState
import numpy as np
import cv2

class RobotTest:
    def __init__(self):
        rospy.init_node('robot_test_node')
        
        # 创建发布者和订阅者
        self.twist_pub = rospy.Publisher('/solve_request_from_vision', 
                                       Twist, queue_size=10)
        self.joint_state_pub = rospy.Publisher('/joint_states', 
                                             JointState, queue_size=10)
        rospy.Subscriber('/joint_angle_result', 
                        Float64MultiArray, self.joint_callback)
        
        # 初始化joint_state消息
        self.joint_state_msg = JointState()
        self.joint_state_msg.name = ['joint_p1', 'joint_r1', 'joint_r2', 
                                    'joint_r3', 'joint_r4', 'joint_r5', 'joint_r6']
        
        self.rate = rospy.Rate(10)  # 10Hz
        
        # 创建滑动条窗口
        cv2.namedWindow('Robot Control')
        
        # 创建滑动条 (窗口名称，滑动条名称，初始值，最大值，回调函数)
        cv2.createTrackbar('X (mm)', 'Robot Control', 0, 1000, lambda x: None)
        cv2.createTrackbar('Y (mm)', 'Robot Control', 0, 1000, lambda x: None)
        cv2.createTrackbar('Z (mm)', 'Robot Control', 0, 1000, lambda x: None)
        cv2.createTrackbar('Roll (deg)', 'Robot Control', 0, 360, lambda x: None)
        cv2.createTrackbar('Pitch (deg)', 'Robot Control', 0, 360, lambda x: None)
        cv2.createTrackbar('Yaw (deg)', 'Robot Control', 0, 360, lambda x: None)
        
        # 设置滑动条初始位置在中间
        cv2.setTrackbarPos('X (mm)', 'Robot Control', 500) 
        cv2.setTrackbarPos('Y (mm)', 'Robot Control', 500)
        cv2.setTrackbarPos('Z (mm)', 'Robot Control', 500)
        cv2.setTrackbarPos('Roll (deg)', 'Robot Control', 180)
        cv2.setTrackbarPos('Pitch (deg)', 'Robot Control', 180)
        cv2.setTrackbarPos('Yaw (deg)', 'Robot Control', 180)

        self.previousState = [0,0,0,0,0,0]
        
    def joint_callback(self, msg):
        # 接收到求解结果后，转换为joint_state并发布
        self.joint_state_msg.header.stamp = rospy.Time.now()
        self.joint_state_msg.position = msg.data
        self.joint_state_msg.velocity = []
        self.joint_state_msg.effort = []
        self.joint_state_pub.publish(self.joint_state_msg)
        
    def send_test_command(self):
        # 获取滑动条位置并转换到正确的范围
        x = (cv2.getTrackbarPos('X (mm)', 'Robot Control') - 500)
        y = (cv2.getTrackbarPos('Y (mm)', 'Robot Control') - 500)
        z = (cv2.getTrackbarPos('Z (mm)', 'Robot Control') - 500)
        
        roll = (cv2.getTrackbarPos('Roll (deg)', 'Robot Control') - 180) * np.pi / 180.0  # -pi到pi
        pitch = (cv2.getTrackbarPos('Pitch (deg)', 'Robot Control') - 180) * np.pi / 180.0
        yaw = (cv2.getTrackbarPos('Yaw (deg)', 'Robot Control') - 180) * np.pi / 180.0
        
        if [x,y,z,roll,pitch,yaw]==self.previousState:
            return 
        # 创建并发布Twist消息
        twist_msg = Twist()
        twist_msg.linear.x = x
        twist_msg.linear.y = y
        twist_msg.linear.z = z
        twist_msg.angular.x = roll
        twist_msg.angular.y = pitch
        twist_msg.angular.z = yaw
        self.twist_pub.publish(twist_msg)

        self.previousState = [x,y,z,roll,pitch,yaw]

        rospy.loginfo(self.previousState)
        
    def run(self):
        # 创建一个显示用的图像
        display = np.zeros((100, 500, 3), dtype=np.uint8)
        cv2.putText(display, "Control Window", (150, 50), 
                   cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
        
        while not rospy.is_shutdown():
            cv2.imshow('Robot Control', display)
            key = cv2.waitKey(1)
            if key == 27:  # ESC键退出
                break
            self.send_test_command()
            self.rate.sleep()
            
        cv2.destroyAllWindows()

if __name__ == '__main__':
    try:
        test_node = RobotTest()
        test_node.run()
    except rospy.ROSInterruptException:
        pass
